# 导入必要的库和模块
from typing import Literal, Optional, Union  # 导入类型提示模块

import torch  # 导入PyTorch库，用于张量操作和深度学习
from torch import nn  # 从PyTorch导入神经网络模块

from .positional_embeddings import PositionEmbeddings  # 导入位置嵌入模块


class CompressiveMemory(nn.Module):
    """实现压缩Transformer的记忆模块，来源于Munkhdalai等人的论文
    "Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention"。
    """

    def __init__(
        self,
        dim_input: int,  # 输入维度
        dim_key: int,  # 键的维度
        dim_value: int,  # 值的维度
        num_heads: int,  # 注意力头的数量
        segment_len: int,  # 段的长度（必须是输入序列长度的因子）
        sampling_factor: Optional[
            int
        ] = None,  # Mixture-of-Depths机制的采样因子的倒数，默认为None
        update: str = "linear",  # 使用的记忆更新规则类型（"linear"或"delta"），默认为"linear"
        causal: bool = False,  # 是否使用因果注意力掩码，默认为False
        position_embedder: Optional[
            PositionEmbeddings
        ] = None,  # 位置嵌入模块，默认为None
        init_state_learnable: bool = False,  # 初始记忆和归一化是否可学习，默认为False
    ):
        """初始化模块。"""
        super(CompressiveMemory, self).__init__()  # 调用父类的初始化方法

        # 记录输入参数
        self.num_heads = num_heads  # 注意力头的数量
        self.segment_len = segment_len  # 段的长度
        self.sampling_factor = sampling_factor  # 采样因子

        self.dim_input = dim_input  # 输入维度
        self.dim_key = dim_key  # 键的维度
        self.dim_value = dim_value  # 值的维度

        self.update = update  # 记忆更新规则

        self.causal = causal  # 是否使用因果注意力掩码

        self.position_embedder = position_embedder  # 位置嵌入模块

        # 为堆叠的SDP注意力定义投影层
        self.proj_k = nn.Linear(
            dim_input, num_heads * dim_key, bias=False
        )  # 键的投影层
        self.proj_v = nn.Linear(
            dim_input, num_heads * dim_value, bias=False
        )  # 值的投影层
        self.proj_q = nn.Linear(
            dim_input, num_heads * dim_key, bias=False
        )  # 查询的投影层

        # 初始化用于计算加权平均的beta值
        self.betas = nn.Parameter(torch.randn(1, num_heads, 1, dim_value))  # beta值

        # 为输出定义投影层
        self.proj_out = nn.Linear(
            num_heads * dim_value, dim_input, bias=False
        )  # 输出投影层

        # 如果init_state_learnable为True，则创建初始记忆矩阵和归一化向量的参数
        if init_state_learnable:
            self.init_mem = nn.Parameter(
                torch.randn(1, self.num_heads, self.dim_key, self.dim_value)
            )  # 初始记忆矩阵
            self.init_z = nn.Parameter(
                torch.ones(1, self.num_heads, self.dim_key, 1)
            )  # 初始归一化向量
        else:
            self.init_mem = None  # 初始记忆矩阵
            self.init_z = None  # 初始归一化向量

    def forward(
        self, x: torch.Tensor, sample_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        """
        应用压缩记忆注意力机制。

        Args:
            x (torch.Tensor): 输入张量，形状为（批大小，序列长度，输入维度）。
            sample_mask (Optional[torch.Tensor], optional): 采样掩码张量，形状为（批大小，序列长度），默认为None。
        Returns:
            torch.Tensor: 输出张量，形状为（批大小，序列长度，输入维度）。
        """
        batch_size, seq_len, _ = x.shape  # 获取批大小、序列长度和输入维度

        num_segments, rem = divmod(
            seq_len, self.segment_len
        )  # 计算段的数量和剩余的序列长度
        num_segments += 1 if rem > 0 else 0  # 如果剩余的序列长度大于0，则增加段的数量

        out = []  # 初始化输出列表

        # 初始化记忆和归一化
        if self.init_mem is not None and self.init_z is not None:
            mem = self.init_mem  # 使用初始记忆矩阵
            z = self.init_z  # 使用初始归一化向量
        else:
            # !!! 论文中没有指定初始化方式，因此这里使用了一个合理的初始化方式
            mem = torch.zeros(
                1, self.num_heads, self.dim_key, self.dim_value
            )  # 初始化记忆矩阵
            z = (
                torch.ones(batch_size, self.num_heads, self.dim_key, 1) / self.dim_key
            )  # 初始化归一化向量

        # 投影输入张量以获取键、值和查询张量
        k_full = (
            self.proj_k(x)
            .unsqueeze(1)
            .view((batch_size, self.num_heads, x.size(1), self.dim_key))  # 键的投影
        )
        v_full = (
            self.proj_v(x)
            .unsqueeze(1)
            .view((batch_size, self.num_heads, x.size(1), self.dim_value))  # 值的投影
        )
        q_full = (
            self.proj_q(x)
            .unsqueeze(1)
            .view((batch_size, self.num_heads, x.size(1), self.dim_key))  # 查询的投影
        )

        for ix in range(num_segments):
            ix_lo = ix * self.segment_len  # 段的起始索引
            ix_hi = min(ix_lo + self.segment_len, x.size(1))  # 段的结束索引
            seg_len = ix_hi - ix_lo  # 段的长度

            # 提取段的键、值和查询张量
            k = k_full[:, :, ix_lo:ix_hi, :]  # 键的提取
            v = v_full[:, :, ix_lo:ix_hi, :]  # 值的提取
            q = q_full[:, :, ix_lo:ix_hi, :]  # 查询的提取

            # 如果提供了采样掩码，则提取段的采样掩码
            if sample_mask is not None:
                if self.sampling_factor is None:
                    raise ValueError(
                        "sampling_factor必须指定，如果sample_mask提供"
                    )  # 如果采样因子未指定，则抛出异常
                ix_lo_seg = (
                    ix * self.segment_len * self.sampling_factor
                )  # 段的起始索引（采样）
                ix_hi_seg = min(
                    ix_lo_seg + self.segment_len * self.sampling_factor,
                    sample_mask.size(1),
                )  # 段的结束索引（采样）
                sample_mask_seg = sample_mask[:, ix_lo_seg:ix_hi_seg]  # 采样掩码的提取
            else:
                sample_mask_seg = None  # 如果未提供采样掩码，则设置为None

            # 如果指定了位置嵌入模块，则添加位置嵌入到查询和键张量中
            if self.position_embedder is not None:
                if sample_mask is None:
                    k_pos = self.position_embedder(
                        k, total_seq_len=seq_len, offset=ix_lo
                    )  # 键的位置嵌入
                    q_pos = self.position_embedder(
                        q, total_seq_len=seq_len, offset=ix_lo
                    )  # 查询的位置嵌入
                else:
                    k_pos = self.position_embedder(
                        k,
                        total_seq_len=seq_len,
                        offset=ix_lo_seg,
                        select_mask=sample_mask_seg,
                    )  # 键的位置嵌入（采样）
                    q_pos = self.position_embedder(
                        q,
                        total_seq_len=seq_len,
                        offset=ix_lo_seg,
                        select_mask=sample_mask_seg,
                    )  # 查询的位置嵌入（采样）

            # 预计算σ(q)用于更新记忆和计算注意力
            # 计算方式见论文第4页的“Memory retrieval”部分
            # 形状：（批大小，注意力头数量，段长度，键维度）
            sigma_q = nn.functional.elu(q) + 1.0  # σ(q)的计算

            # 应用SDP注意力，见论文公式（2）
            if self.position_embedder is not None:
                scores = (
                    q_pos @ k_pos.transpose(-2, -1) / self.dim_key**0.5
                )  # 注意力得分的计算（位置嵌入）
            else:
                scores = q @ k.transpose(-2, -1) / self.dim_key**0.5  # 注意力得分的计算

            # 如果指定了因果掩码，则计算和应用因果掩码
            if self.causal:
                mask = torch.tril(
                    torch.ones((seg_len, seg_len), dtype=torch.bool), diagonal=0
                )  # 因果掩码的计算
                mask = (
                    mask.unsqueeze(0)
                    .unsqueeze(0)
                    .repeat((batch_size, self.num_heads, 1, 1))
                )  # 因果掩码的重复
                scores.masked_fill_(
                    torch.logical_not(mask), float("-inf")
                )  # 因果掩码的应用

            # 计算SDP注意力，见论文公式（2）
            att_dot = nn.functional.softmax(scores, dim=-1) @ v  # SDP注意力的计算

            # 计算归一化的线性注意力，见论文公式（3）
            # 形状：（批大小，注意力头数量，段长度，值维度）
            att_mem = (sigma_q @ mem) / (sigma_q @ z)  # 归一化的线性注意力的计算

            # 更新记忆，见论文公式（4）和（5）
            sigma_k = nn.functional.elu(k) + 1.0  # σ(k)的计算
            if self.update == "linear":
                mem = mem + sigma_k.transpose(-2, -1) @ v  # 线性更新
            elif self.update == "delta":
                mem = mem + sigma_k.transpose(-2, -1) @ (
                    v - (sigma_k @ mem) / (sigma_k @ z)
                )  # Delta更新

            # 更新归一化项，见论文公式（4）
            z = z + (nn.functional.elu(k) + 1.0).sum(dim=-2, keepdim=True).transpose(
                -2, -1
            )  # 归一化项的更新

            # 计算加权平均的注意力，见论文公式（6）
            att = (
                nn.functional.sigmoid(self.betas) * att_mem
                + (1 - nn.functional.sigmoid(self.betas)) * att_dot
            )  # 加权平均的注意力的计算
            att = att.view(
                (batch_size, seg_len, self.num_heads * self.dim_value)
            )  # 注意力的形状调整

            # 将输出追加到缓冲区
            # 计算方式见论文公式（7）
            out.append(self.proj_out(att))  # 输出的追加

        # 返回连接的全序列输出
        out = torch.concat(out, dim=1)  # 输出的连接

        return out


def test_compressive_memory(
    short_seq_len: bool = False,
    even_seq_len: bool = True,
    causal_masking: bool = False,
    update: str = "linear",
) -> None:
    # 设置示例模块参数
    dim_input = 512  # 输入维度
    dim_key = 64  # 键的维度
    dim_value = 64  # 值的维度
    num_heads = 8  # 注意力头的数量
    segment_len = 32  # 段的长度
    causal = causal_masking  # 是否使用因果注意力掩码

    # 设置示例输入维度
    batch_size = 4  # 批大小

    # 处理序列长度
    if short_seq_len:
        seq_len = 16  # 短序列长度
    else:
        if even_seq_len:
            seq_len = 128  # 偶序列长度
        else:
            seq_len = 144  # 奇序列长度

    # 初始化模块
    model = CompressiveMemory(
        dim_input, dim_key, dim_value, num_heads, segment_len, update, causal
    )

    # 生成随机输入
    batch = torch.randn(batch_size, seq_len, dim_input)  # 随机输入

    # 应用压缩记忆模块
    model(batch)


if __name__ == "__main__":
    # 测试所有情况
    print("测试所有情况：")

    short_seq_len = True  # 短序列长度
    # 在这种情况下，even_seq_len不重要，因此任意设置为True
    even_seq_len = True

    for causal_masking in [True, False]:
        for update in ["linear", "delta"]:
            print(f"  测试causal_masking={causal_masking}和update={update}")
            test_compressive_memory(
                short_seq_len=short_seq_len,
                even_seq_len=even_seq_len,
                causal_masking=causal_masking,
                update=update,
            )

    # 测试所有情况
    print("测试所有情况：")

    short_seq_len = False  # 非短序列长度

    for even_seq_len in [True, False]:
        for causal_masking in [True, False]:
            for update in ["linear", "delta"]:
                print(
                    f"  测试even_seq_len={even_seq_len}, causal_masking={causal_masking}和update={update}"
                )
                test_compressive_memory(
                    short_seq_len=short_seq_len,
                    even_seq_len=even_seq_len,
                    causal_masking=causal_masking,
                    update=update,
                )
